"""
数据库工具模块
"""
from sqlalchemy import create_engine, text as sa_text
from sqlalchemy.pool import QueuePool
from app.config.config import DevelopmentConfig
from app.utils.log_control import INFO, ERROR
import time
import urllib.parse
from contextlib import contextmanager
from sqlalchemy.engine import Row


class DatabaseConnection:
    # 创建数据库引擎
    password = urllib.parse.quote_plus(DevelopmentConfig.MYSQL_PASSWORD)  # 处理特殊字符
    engine = create_engine(
        f'mysql+pymysql://{DevelopmentConfig.MYSQL_USER}:{password}'
        f'@{DevelopmentConfig.MYSQL_HOST}:{DevelopmentConfig.MYSQL_PORT}/{DevelopmentConfig.MYSQL_DB}',
        pool_size=20,  # 增加连接池大小
        max_overflow=30,  # 增加最大溢出连接数
        pool_timeout=30,  # 减少连接池超时时间
        pool_recycle=1800,  # 保持连接回收时间
        pool_pre_ping=True,  # 保持连接预检
        echo=False,
        future=True,
        connect_args={
            'connect_timeout': 30,  # 减少连接超时时间
            'read_timeout': 30,
            'write_timeout': 30,
            'program_name': 'TDS_APP',
            'charset': 'utf8mb4',
            'autocommit': True,  # 启用自动提交
            'client_flag': 2,
            'ssl': None,
            'init_command': 'SET SESSION wait_timeout=1800, interactive_timeout=1800, sql_mode="STRICT_TRANS_TABLES,NO_ZERO_IN_DATE,NO_ZERO_DATE,ERROR_FOR_DIVISION_BY_ZERO,NO_ENGINE_SUBSTITUTION"'
        }
    )

    def __init__(self):
        self.connection = None
        self._in_transaction = False

    @contextmanager
    def get_connection(self):
        """获取数据库连接的上下文管理器"""
        retry_count = 3
        retry_delay = 1.0
        last_error = None

        while retry_count > 0:
            try:
                if self.connection is not None:
                    try:
                        self.connection.close()
                    except Exception as e:
                        INFO.logger.warning(f"关闭旧连接时出现警告: {str(e)}")

                connection = self.engine.connect()
                try:
                    yield connection
                finally:
                    if connection is not None:
                        connection.close()
                return
            except Exception as e:
                last_error = e
                retry_count -= 1
                if retry_count > 0:
                    time.sleep(retry_delay)
                    retry_delay *= 2
                else:
                    ERROR.logger.error(f"获取数据库连接失败: {str(e)}")
                    raise last_error

    def execute(self, sql, params=None):
        """执行更新SQL"""
        retry_count = 3
        retry_delay = 1.0
        last_error = None
        success = False

        while retry_count > 0 and not success:
            try:
                with self.get_connection() as connection:
                    if isinstance(sql, str):
                        sql = sa_text(sql)
                    
                    trans = connection.begin()
                    try:
                        result = connection.execute(sql, params or {})
                        trans.commit()
                        success = True
                        return result.rowcount
                    except Exception as e:
                        trans.rollback()
                        ERROR.logger.error(f"执行SQL失败，已回滚: {str(e)}")
                        raise
            except Exception as e:
                last_error = e
                retry_count -= 1
                if retry_count > 0:
                    ERROR.logger.warning(f"SQL执行失败，准备重试: {str(e)}")
                    time.sleep(retry_delay)
                    retry_delay *= 2
                    continue
                ERROR.logger.error(f"SQL执行失败: {str(e)}, SQL: {sql}")
                raise last_error

    def execute_query(self, sql, params=None):
        """执行查询SQL"""
        start_time = time.time()
        retry_count = 3
        retry_delay = 1.0
        last_error = None

        while retry_count > 0:
            try:
                with self.get_connection() as connection:
                    if isinstance(sql, str):
                        sql = sa_text(sql)
                    
                    with connection.begin():
                        result = connection.execute(sql, params or {})
                        # 将 Row 对象转换为字典
                        data = [dict(row) for row in result.mappings()]
                        
                        query_time = round((time.time() - start_time) * 1000)
                        return {
                            'data': data,
                            'query_time': query_time
                        }
            except Exception as e:
                last_error = e
                retry_count -= 1
                if retry_count > 0:
                    time.sleep(retry_delay)
                    retry_delay *= 2
                    continue
                ERROR.logger.error(f"SQL执行失败: {str(e)}, SQL: {sql}")
                raise last_error

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is not None:
            ERROR.logger.error(f"数据库操作出错: {str(exc_val)}")

